import cv2
import numpy as np
import yaml
import os
from datetime import datetime

class CircleCapture:
    def __init__(self, camera_config_path='camera_calibration.yaml'):
        """Initialize camera and parameters with proper calibration"""
        self.camera_matrix = None
        self.dist_coeffs = None
        self.map_x = None
        self.map_y = None
        self.roi = None
        self.load_camera_calibration(camera_config_path)
        
        self.cap = cv2.VideoCapture(0)
        if not self.cap.isOpened():
            raise Exception("Cannot open camera")
        
        # Set camera parameters
        cap_width, cap_height = 1280,720
        #cap_width, cap_height = 640, 480
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, cap_width)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, cap_height)
        
        # Get actual resolution
        ret, test_frame = self.cap.read()
        if not ret:
            raise Exception("Cannot read from camera")
        
        self.resolution = (test_frame.shape[1], test_frame.shape[0])
        print(f"Camera resolution: {self.resolution[0]}x{self.resolution[1]}")
        
        # Setup undistortion maps for better performance
        self.setup_undistort_maps()
        
        self.circle_pixel_coords = []  # Store circle center pixel coordinates
        self.circle_world_coords = []  # Store circle center world coordinates
        self.current_image = None
        self.show_undistorted = True  # Flag to toggle between original and undistorted view
        
        # 用于半径平滑的历史记录
        self.radius_history = {}  # key: (x,y)中心坐标, value: list of radii
        self.max_history = 10     # 保存的历史帧数
        self.smoothing_weight = 0.7  # 新测量的权重
    
    def load_camera_calibration(self, config_path):
        """Load camera calibration parameters"""
        if not os.path.exists(config_path):
            print(f"Warning: Camera calibration file {config_path} does not exist")
            # Create default identity matrices for testing
            self.camera_matrix = np.eye(3)
            self.dist_coeffs = np.zeros((5, 1))
            print("Using default identity matrices for testing")
            return
            
        try:
            with open(config_path, 'r') as file:
                calibration_data = yaml.safe_load(file)
            
            # Extract calibration parameters
            if 'camera_matrix' in calibration_data:
                self.camera_matrix = np.array(calibration_data['camera_matrix'])
            elif 'mtx' in calibration_data:
                self.camera_matrix = np.array(calibration_data['mtx'])
            else:
                raise KeyError("Could not find camera matrix in calibration file")
            
            if 'dist_coeffs' in calibration_data:
                self.dist_coeffs = np.array(calibration_data['dist_coeffs'])
            elif 'dist' in calibration_data:
                self.dist_coeffs = np.array(calibration_data['dist'])
            elif 'distortion_coefficients' in calibration_data:
                self.dist_coeffs = np.array(calibration_data['distortion_coefficients'])
            else:
                self.dist_coeffs = np.zeros((5, 1))
                print("Warning: Using zero distortion coefficients")
            
            # Print calibration info if available
            if 'calibration_date' in calibration_data:
                print(f"Calibration date: {calibration_data['calibration_date']}")
            if 'reprojection_error' in calibration_data:
                print(f"Reprojection error: {calibration_data['reprojection_error']:.6f} pixels")
            
            print("Camera calibration parameters loaded successfully")
            print(f"Camera matrix:\n{self.camera_matrix}")
            print(f"Distortion coefficients: {self.dist_coeffs.flatten()}")
            
        except Exception as e:
            print(f"Error loading camera calibration: {e}")
            # Fallback to identity matrices
            self.camera_matrix = np.eye(3)
            self.dist_coeffs = np.zeros((5, 1))
            print("Using default identity matrices due to loading error")
    
    def setup_undistort_maps(self):
        """Setup undistortion maps for better performance"""
        if self.camera_matrix is not None and self.dist_coeffs is not None:
            try:
                # Get optimal new camera matrix and ROI
                new_camera_matrix, roi = cv2.getOptimalNewCameraMatrix(
                    self.camera_matrix, self.dist_coeffs, self.resolution, 1, self.resolution
                )
                self.roi = roi
                
                # Generate undistortion maps
                self.map_x, self.map_y = cv2.initUndistortRectifyMap(
                    self.camera_matrix, self.dist_coeffs, None, 
                    new_camera_matrix, self.resolution, cv2.CV_32FC1
                )
                print("Undistortion maps created successfully")
                
            except Exception as e:
                print(f"Error creating undistortion maps: {e}")
                self.map_x, self.map_y = None, None
        else:
            self.map_x, self.map_y = None, None
    
    def undistort_image(self, frame):
        """Undistort image using precomputed maps for better performance"""
        if self.map_x is not None and self.map_y is not None:
            try:
                undistorted = cv2.remap(frame, self.map_x, self.map_y, cv2.INTER_LINEAR)
                
                # Crop the image to the ROI if available
                if self.roi is not None:
                    x, y, w, h = self.roi
                    if w > 0 and h > 0:
                        undistorted = undistorted[y:y+h, x:x+w]
                
                return undistorted
            except Exception as e:
                print(f"Error in undistortion: {e}")
                return frame
        return frame
    
    def fit_circle_least_squares(self, contour):
        """最小二乘圆拟合"""
        pts = contour.reshape(-1, 2).astype(np.float32)
        M = np.zeros((len(pts), 3))
        F = np.zeros(len(pts))
        
        for i, (x, y) in enumerate(pts):
            M[i] = [x, y, 1]
            F[i] = -(x*x + y*y)
            
        # 求解线性方程
        A = M.T.dot(M)
        B = M.T.dot(F)
        
        try:
            X = np.linalg.solve(A, B)
        except np.linalg.LinAlgError:
            # 退化情况，使用简单方法
            mean_pt = np.mean(pts, axis=0)
            dists = np.linalg.norm(pts - mean_pt, axis=1)
            return mean_pt, np.mean(dists)
            
        # 计算圆心和半径
        cx, cy = -X[0]/2, -X[1]/2
        r = np.sqrt(cx*cx + cy*cy - X[2])
        
        return (cx, cy), r
    
    def detect_circles(self, frame):
        """Detect circles with contour detection and least squares fitting"""
        # 1. 预处理
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        blurred = cv2.GaussianBlur(gray, (9, 9), 2)
        
        # 自适应阈值处理
        binary = cv2.adaptiveThreshold(
            blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY_INV, 11, 2
        )
        
        # 2. 轮廓检测
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        detected_circles = []
        for cnt in contours:
            # 过滤小轮廓
            area = cv2.contourArea(cnt)
            if area < 50 or area > 50000:
                continue
                
            # 周长和面积判断圆度
            perimeter = cv2.arcLength(cnt, True)
            if perimeter == 0:
                continue
                
            circularity = 4 * np.pi * area / (perimeter * perimeter)
            if not (0.7 <= circularity <= 1.2):
                continue
                
            # 3. 最小二乘圆拟合
            center, radius = self.fit_circle_least_squares(cnt)
            
            # 验证半径范围
            if 10 <= radius <= 100:
                detected_circles.append((int(center[0]), int(center[1]), int(radius)))
                
        return detected_circles
    
    def run_capture(self):
        """Run image capture program with improved interface"""
        print("=== Circle Center Capture with Camera Calibration ===")
        print("按键说明:")
        print("  SPACE - 保存当前图像和圆心坐标")
        print("  'q'   - 退出程序")
        print("  'u'   - 切换原始/校正视图")
        print("  'c'   - 清除所有已保存的点")
        print("  's'   - 保存数据到文件")
        print(f"已保存点数: {len(self.circle_pixel_coords)}")
        print("请确保图像中有清晰的圆形标记")
        
        while True:
            ret, frame = self.cap.read()
            if not ret:
                print("Failed to capture frame")
                break
                
            # Process frame based on view mode
            if self.show_undistorted:
                processed_frame = self.undistort_image(frame)
                view_label = "Undistorted View"
            else:
                processed_frame = frame.copy()
                view_label = "Original View"
                
            self.current_image = processed_frame.copy()
            
            # Detect circles
            circles = self.detect_circles(processed_frame)
            
            # 应用时间滤波平滑半径
            smoothed_circles = []
            for (x, y, r) in circles:
                # 找到最接近的历史圆心
                key = None
                min_dist = float('inf')
                for hist_key in list(self.radius_history.keys()):
                    dist = np.hypot(x - hist_key[0], y - hist_key[1])
                    if dist < 10 and dist < min_dist:  # 阈值可根据实际情况调整
                        min_dist = dist
                        key = hist_key
                        
                if key is not None:
                    # 更新历史记录
                    self.radius_history[key].append(r)
                    if len(self.radius_history[key]) > self.max_history:
                        self.radius_history[key].pop(0)
                        
                    # 计算平滑半径（指数移动平均）
                    smoothed_r = r
                    for prev_r in reversed(self.radius_history[key][:-1]):
                        smoothed_r = self.smoothing_weight * smoothed_r + (1 - self.smoothing_weight) * prev_r
                        
                    smoothed_circles.append((x, y, int(round(smoothed_r))))
                    
                    # 更新键为新坐标
                    self.radius_history[(x, y)] = self.radius_history.pop(key)
                else:
                    # 新检测到的圆
                    self.radius_history[(x, y)] = [r]
                    smoothed_circles.append((x, y, r))
            
            circles = smoothed_circles
            
            # Draw detection results
            display_frame = processed_frame.copy()
            for (x, y, r) in circles:
                # Draw circle
                cv2.circle(display_frame, (x, y), r, (0, 255, 0), 1)
                # Draw circle center
                cv2.circle(display_frame, (x, y), 2, (0, 0, 255), 2)
                # Display coordinates and radius
                cv2.putText(display_frame, f"({x}, {y}), r={r}", (x + 10, y - 10), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
            
            # Add informational text
            cv2.putText(display_frame, f"Detected: {len(circles)} circles", (10, 30), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
            cv2.putText(display_frame, f"Saved: {len(self.circle_pixel_coords)} points", (10, 60), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 0), 1)
            cv2.putText(display_frame, view_label, (10, 90), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 255), 1)
            
            # Add control instructions
            instructions = [
                "SPACE: Save circles",
                "q: Quit",
                "u: Toggle view",
                "c: Clear points",
                "s: Save to file"
            ]
            
            for i, instruction in enumerate(instructions):
                cv2.putText(display_frame, instruction, (10, 120 + i*25), 
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (200, 200, 200), 1)
            
            cv2.imshow('Circle Detection - Press SPACE to save', display_frame)
            
            key = cv2.waitKey(1) & 0xFF
            if key == ord(' '):  # Space key to save
                if len(circles) > 0:
                    self.save_circle_data(circles)
                else:
                    print("No circles detected, please adjust the image")
                    
            elif key == ord('q'):  # q key to quit
                break
                
            elif key == ord('u'):  # u key to toggle view
                self.show_undistorted = not self.show_undistorted
                print(f"View mode: {'Undistorted' if self.show_undistorted else 'Original'}")
                
            elif key == ord('c'):  # c key to clear points
                self.circle_pixel_coords.clear()
                self.circle_world_coords.clear()
                self.radius_history.clear()  # 清除半径历史
                print("All saved points and radius history cleared")
                
            elif key == ord('s'):  # s key to save data
                self.save_to_file()
        
        self.cap.release()
        cv2.destroyAllWindows()
        
    def save_circle_data(self, circles):
        """Save circle center data with improved user interaction"""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # Save image
        image_filename = f"calibration_image_{timestamp}.jpg"
        cv2.imwrite(image_filename, self.current_image)
        print(f"Image saved: {image_filename}")
        
        # Save pixel coordinates and radius
        pixel_coords = [(x, y, r) for (x, y, r) in circles]
        self.circle_pixel_coords.extend(pixel_coords)
        
        # Prompt user to input world coordinates
        print(f"\nDetected {len(circles)} circles, please enter corresponding world coordinates (X, Y):")
        world_coords_batch = []
        for i, (x, y, r) in enumerate(circles):
            try:
                print(f"Circle {i+1} at pixel position ({x}, {y}), radius {r}")
                world_x = float(input("  World coordinate X: "))
                world_y = float(input("  World coordinate Y: "))
                world_coords_batch.append((world_x, world_y))
                print(f"  Saved: ({world_x}, {world_y})")
            except ValueError:
                print("Invalid input, please enter numbers. This circle will be skipped.")
                # Remove the corresponding pixel coordinate
                if len(self.circle_pixel_coords) >= len(pixel_coords):
                    self.circle_pixel_coords.pop()
        
        self.circle_world_coords.extend(world_coords_batch)
        
        # Save to file
        self.save_to_file()
        print(f"Data saved. Total points: {len(self.circle_pixel_coords)}")
        
    def save_to_file(self):
        """Save data to file with improved formatting"""
        if len(self.circle_pixel_coords) != len(self.circle_world_coords):
            print("Error: Number of pixel coordinates and world coordinates do not match")
            return
            
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"calibration_data_{timestamp}.txt"
        
        with open(filename, 'w') as f:
            # Write header with calibration info
            f.write("# Camera calibration data - Circle centers and radii\n")
            f.write(f"# Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write(f"# Total points: {len(self.circle_pixel_coords)}\n")
            f.write("# Format: pixel_x, pixel_y, radius, world_x, world_y\n")
            f.write("pixel_x,pixel_y,radius,world_x,world_y\n")
            
            for (px, py, r), (wx, wy) in zip(self.circle_pixel_coords, self.circle_world_coords):
                f.write(f"{px},{py},{r},{wx},{wy}\n")
                
        print(f"Saved {len(self.circle_pixel_coords)} calibration points to {filename}")

if __name__ == "__main__":
    try:
        capture = CircleCapture()
        capture.run_capture()
    except Exception as e:
        print(f"Program error: {e}")
